load("@local_config_cuda//cuda:build_defs.bzl", "cuda_library")
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("//xla:xla.default.bzl", "xla_cc_test", "xla_internal")
load("//xla/tests:build_defs.bzl", "xla_test")
load("//xla/tsl:tsl.bzl", "internal_visibility", "nvtx_headers")
load("//xla/tsl/platform:build_config.bzl", "tf_proto_library")
load("//xla/tsl/platform/default:cuda_build_defs.bzl", "if_cuda_is_configured")

package(
    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
    default_visibility = internal_visibility([":friends"]),
    licenses = ["notice"],
)

package_group(
    name = "friends",
    includes = ["//xla:friends"],
)

#===-------------------------------------------------------------------------------------------===//
# Runtime tracing libraries
#===-------------------------------------------------------------------------------------------===//

cc_library(
    name = "thunk_runtime_dependencies",
    # Register GPU collectives plugin.
    deps = ["//xla/backends/gpu/collectives:gpu_collectives_plugin"],
)

cc_library(
    name = "annotation",
    srcs = ["annotation.cc"],
    hdrs = ["annotation.h"],
    local_defines = if_cuda_is_configured(["GOOGLE_CUDA=1"]),
    deps = [
        "//xla:printer",
        "//xla/hlo/ir:hlo",
        "//xla/hlo/utils:hlo_longest_prefix",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
        "@tsl//tsl/platform:errors",
        "@tsl//tsl/profiler/lib:nvtx_utils",
        "@tsl//tsl/profiler/lib:scoped_annotation",
    ] + if_cuda_is_configured(nvtx_headers()),
)

#===-------------------------------------------------------------------------------------------===//
# Command Buffer Integration
#===-------------------------------------------------------------------------------------------===//

cc_library(
    name = "command_buffer_cmd",
    srcs = ["command_buffer_cmd.cc"],
    hdrs = ["command_buffer_cmd.h"],
    deps = [
        ":all_gather_thunk",
        ":all_reduce_thunk",
        ":all_to_all_thunk",
        ":annotation",
        ":collective_broadcast_thunk",
        ":collective_thunk",
        ":copy_thunk",
        ":custom_call_thunk",
        ":dynamic_slice_thunk",
        ":gpublas_lt_matmul_thunk",
        ":thunk",
        ":while_thunk",
        "//xla:debug_options_flags",
        "//xla:executable_run_options",
        "//xla:shape_util",
        "//xla:status_macros",
        "//xla:types",
        "//xla:util",
        "//xla/backends/gpu/collectives:gpu_clique_key",
        "//xla/backends/gpu/collectives:gpu_collectives",
        "//xla/ffi:call_frame",
        "//xla/ffi:ffi_api",
        "//xla/ffi/api:c_api",
        "//xla/hlo/ir:hlo",
        "//xla/runtime:buffer_use",
        "//xla/runtime:execution_graph",
        "//xla/runtime:object_pool",
        "//xla/runtime:resource_use",
        "//xla/service:buffer_assignment",
        "//xla/service:collective_ops_utils",
        "//xla/service:computation_placer",
        "//xla/service:custom_call_status_internal",
        "//xla/service:custom_call_status_public_headers",
        "//xla/service:global_device_id",
        "//xla/service/gpu:buffer_allocations",
        "//xla/service/gpu:launch_dimensions",
        "//xla/service/gpu:matmul_utils",
        "//xla/service/gpu:stream_executor_util",
        "//xla/service/gpu/kernels:custom_kernel",
        "//xla/stream_executor:command_buffer",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:dnn",
        "//xla/stream_executor:kernel",
        "//xla/stream_executor:memory_allocation",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/stream_executor:trace_command_buffer_factory",
        "//xla/stream_executor/gpu:tma_metadata",
        "//xla/tsl/lib/gtl:int_type",
        "//xla/tsl/platform:env",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:btree",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:flat_hash_set",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/functional:function_ref",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
        "@llvm-project//llvm:Support",
        "@tsl//tsl/profiler/lib:scoped_annotation",
    ],
)

xla_test(
    name = "command_buffer_cmd_test",
    srcs = ["command_buffer_cmd_test.cc"],
    backends = ["gpu"],
    deps = [
        ":command_buffer_cmd",
        ":thunk",
        "//xla/runtime:buffer_use",
        "//xla/runtime:resource_use",
        "//xla/service:buffer_assignment",
        "//xla/service:executable",
        "//xla/service:platform_util",
        "//xla/service/gpu:buffer_allocations",
        "//xla/service/gpu:launch_dimensions",
        "//xla/stream_executor:command_buffer",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:platform_manager",
        "//xla/stream_executor:stream_executor_h",
        "//xla/stream_executor:stream_executor_memory_allocator",
        "//xla/stream_executor/gpu:gpu_test_kernels_fatbin",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:status",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/platform:test",
        "//xla/tsl/platform:test_benchmark",
        "//xla/tsl/platform:test_main",
        "//xla/tsl/util:safe_reinterpret_cast",
        "@com_google_absl//absl/functional:function_ref",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/types:span",
    ],
)

cc_library(
    name = "command_buffer_cmd_emitter",
    srcs = ["command_buffer_cmd_emitter.cc"],
    hdrs = ["command_buffer_cmd_emitter.h"],
    deps = [
        ":all_gather_thunk",
        ":all_reduce_thunk",
        ":all_to_all_thunk",
        ":collective_thunk",
        ":command_buffer_cmd",
        ":conditional_thunk",
        ":copy_thunk",
        ":cudnn_thunk",
        ":custom_call_thunk",
        ":dynamic_slice_thunk",
        ":gemm_thunk",
        ":gpublas_lt_matmul_thunk",
        ":kernel_thunk",
        ":memset_thunk",
        ":replica_id_thunk",
        ":sequential_thunk",
        ":thunk",
        ":while_thunk",
        "//xla:util",
        "//xla/runtime:buffer_use",
        "//xla/runtime:resource_use",
        "//xla/service:buffer_assignment",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
    ],
)

#===-------------------------------------------------------------------------------------------===//
# XLA Thunks Runtime
#===-------------------------------------------------------------------------------------------===//

cc_library(
    name = "dynamic_slice_thunk",
    srcs = ["dynamic_slice_thunk.cc"],
    hdrs = ["dynamic_slice_thunk.h"],
    deps = [
        ":sequential_thunk",
        ":thunk",
        "//xla:literal",
        "//xla:literal_util",
        "//xla:shape_util",
        "//xla:status_macros",
        "//xla/hlo/evaluator:hlo_evaluator",
        "//xla/service:buffer_assignment",
        "//xla/service/gpu:buffer_allocations",
        "//xla/service/gpu:ir_emission_utils",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:memory_allocation",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/functional:function_ref",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/synchronization",
        "@llvm-project//llvm:Support",
        "@tsl//tsl/platform:errors",
        "@tsl//tsl/platform:logging",
        "@tsl//tsl/platform:statusor",
    ],
)

xla_test(
    name = "dynamic_slice_thunk_test",
    srcs = ["dynamic_slice_thunk_test.cc"],
    backends = [
        "gpu",
    ],
    deps = [
        ":custom_call_thunk",
        ":dynamic_slice_thunk",
        ":gemm_thunk",
        ":thunk",
        "//xla:shape_util",
        "//xla:xla_data_proto_cc",
        "//xla/ffi",
        "//xla/ffi:ffi_api",
        "//xla/hlo/ir:hlo",
        "//xla/hlo/testlib:hlo_hardware_independent_test_base",
        "//xla/service:buffer_assignment",
        "//xla/service:executable",
        "//xla/service:platform_util",
        "//xla/service/gpu:buffer_allocations",
        "//xla/service/gpu:matmul_utils",
        "//xla/service/gpu:resource_requests",
        "//xla/stream_executor:blas",
        "//xla/stream_executor:command_buffer",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:device_memory_allocator",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:platform_manager",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/stream_executor:stream_executor_memory_allocator",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "cholesky_thunk",
    srcs = ["cholesky_thunk.cc"],
    hdrs = ["cholesky_thunk.h"],
    deps = [
        ":make_batch_pointers",
        ":thunk",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/service:buffer_assignment",
        "//xla/stream_executor:blas",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:gpu_solver_context",
        "//xla/stream_executor:stream",
        "//xla/tsl/platform:errors",
        "@com_google_absl//absl/functional:any_invocable",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
    ],
)

cc_library(
    name = "command_buffer_thunk",
    srcs = ["command_buffer_thunk.cc"],
    hdrs = ["command_buffer_thunk.h"],
    deps = [
        ":command_buffer_cmd",
        ":sequential_thunk",  # build_cleaner: keep
        ":thunk",
        "//xla/service:buffer_assignment",  # build_cleaner: keep
        "//xla/service/gpu:buffer_allocations",  # build_cleaner: keep
        "//xla/stream_executor:command_buffer",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/platform:env",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:logging",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/functional:function_ref",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/synchronization",
        "@tsl//tsl/profiler/lib:profiler_lock",
        "@tsl//tsl/profiler/lib:traceme",
        "@tsl//tsl/profiler/lib:traceme_encode",
    ] + xla_internal(["service:execution_graph_visualizer_google"]),
)

xla_test(
    name = "command_buffer_thunk_test",
    srcs = ["command_buffer_thunk_test.cc"],
    backends = ["gpu"],
    tags = ["test_migrated_to_hlo_runner_pjrt"],
    deps = [
        ":command_buffer_cmd",
        ":command_buffer_thunk",
        ":dynamic_slice_thunk",
        ":gpublas_lt_matmul_thunk",
        ":memset_thunk",
        ":sequential_thunk",
        ":thunk",
        "//xla:error_spec",
        "//xla:shape_util",
        "//xla:xla_data_proto_cc",
        "//xla:xla_proto_cc",
        "//xla/hlo/ir:hlo",
        "//xla/runtime:buffer_use",
        "//xla/service:buffer_assignment",
        "//xla/service:executable",
        "//xla/service:hlo_module_config",
        "//xla/service:platform_util",
        "//xla/service/gpu:buffer_allocations",
        "//xla/service/gpu:launch_dimensions",
        "//xla/service/gpu:matmul_utils",
        "//xla/service/gpu/kernels:custom_kernel",
        "//xla/stream_executor:blas",
        "//xla/stream_executor:command_buffer",
        "//xla/stream_executor:device_description",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:device_memory_allocator",
        "//xla/stream_executor:kernel",
        "//xla/stream_executor:kernel_spec",
        "//xla/stream_executor:launch_dim",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:platform_manager",
        "//xla/stream_executor:semantic_version",
        "//xla/stream_executor:stream_executor_h",
        "//xla/stream_executor:stream_executor_memory_allocator",
        "//xla/stream_executor/cuda:cuda_compute_capability",
        "//xla/stream_executor/gpu:gpu_test_kernels",
        "//xla/stream_executor/gpu:gpu_test_kernels_fatbin",
        "//xla/tests:hlo_pjrt_interpreter_reference_mixin",
        "//xla/tests:hlo_pjrt_test_base",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/types:span",
        "@com_google_googletest//:gtest_main",
        "@tsl//tsl/profiler/lib:profiler_lock",
    ],
)

xla_test(
    name = "cuda_command_buffer_thunk_test",
    srcs = ["cuda_command_buffer_thunk_test.cc"],
    backends = ["gpu"],
    tags = ["cuda-only"],
    deps = [
        ":command_buffer_cmd",
        ":command_buffer_thunk",
        ":thunk",
        "//xla:types",
        "//xla:xla_data_proto_cc",
        "//xla/runtime:buffer_use",
        "//xla/service:buffer_assignment",
        "//xla/service:executable",
        "//xla/service:platform_util",
        "//xla/service/gpu:buffer_allocations",
        "//xla/stream_executor:command_buffer",
        "//xla/stream_executor:device_description",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:device_memory_allocator",
        "//xla/stream_executor:dnn",
        "//xla/stream_executor:kernel_spec",
        "//xla/stream_executor:numeric_options",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:platform_manager",
        "//xla/stream_executor:stream_executor_h",
        "//xla/stream_executor:stream_executor_memory_allocator",
        "//xla/stream_executor/cuda:cuda_compute_capability",
        "//xla/stream_executor/cuda:cudnn_plugin",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/status:status_matchers",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/types:span",
        "@com_google_googletest//:gtest_main",
        "@cudnn_frontend_archive//:cudnn_frontend",
        "@local_config_cuda//cuda:cuda_headers",
    ],
)

cc_library(
    name = "conditional_thunk",
    srcs = ["conditional_thunk.cc"],
    hdrs = ["conditional_thunk.h"],
    deps = [
        ":host_memory_pool",
        ":sequential_thunk",
        ":thunk",
        ":thunk_proto_cc",
        "//xla:status_macros",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/service:buffer_assignment",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/functional:function_ref",
        "@com_google_absl//absl/functional:overload",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
    ],
)

xla_cc_test(
    name = "conditional_thunk_test",
    srcs = ["conditional_thunk_test.cc"],
    deps = [
        ":conditional_thunk",
        ":sequential_thunk",
        ":thunk",
        ":thunk_proto_cc",
        "//xla/service:buffer_assignment",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/platform:test",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_googletest//:gtest_main",
        "@tsl//tsl/platform:protobuf",
    ],
)

cc_library(
    name = "convolution_thunk",
    srcs = ["convolution_thunk.cc"],
    hdrs = ["convolution_thunk.h"],
    deps = [
        ":thunk",
        "//xla:util",
        "//xla/service:buffer_assignment",
        "//xla/service/gpu:gpu_conv_runner",
        "//xla/service/gpu:stream_executor_util",
        "//xla/stream_executor:device_description",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:dnn",
        "//xla/stream_executor:scratch_allocator",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/protobuf:dnn_proto_cc",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
        "@tsl//tsl/platform:errors",
    ],
)

cc_library(
    name = "copy_thunk",
    srcs = ["copy_thunk.cc"],
    hdrs = ["copy_thunk.h"],
    deps = [
        ":thunk",
        ":thunk_proto_cc",
        ":while_thunk",
        "//xla/hlo/ir:hlo",
        "//xla/service:buffer_assignment",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:event",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/container:node_hash_map",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
    ],
)

xla_cc_test(
    name = "copy_thunk_test",
    srcs = ["copy_thunk_test.cc"],
    deps = [
        ":copy_thunk",
        ":thunk",
        ":thunk_proto_cc",
        "//xla/service:buffer_assignment",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_absl//absl/log:check",
        "@com_google_googletest//:gtest_main",
        "@tsl//tsl/platform:protobuf",
    ],
)

cc_library(
    name = "cub_sort_thunk",
    srcs = ["cub_sort_thunk.cc"],
    hdrs = ["cub_sort_thunk.h"],
    deps = [
        ":thunk",
        "//xla:executable_run_options",
        "//xla:shape_util",
        "//xla:xla_data_proto_cc",
        "//xla/ffi:call_frame",
        "//xla/ffi:ffi_api",
        "//xla/ffi/api:c_api",
        "//xla/service:buffer_assignment",
        "//xla/service/gpu:buffer_allocations",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:device_memory_allocator",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:stream",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/strings:string_view",
    ],
)

cc_library(
    name = "custom_call_thunk",
    srcs = ["custom_call_thunk.cc"],
    hdrs = ["custom_call_thunk.h"],
    deps = [
        ":thunk",
        "//xla:executable_run_options",
        "//xla:shape_util",
        "//xla:util",
        "//xla/ffi:call_frame",
        "//xla/ffi:execution_context",
        "//xla/ffi:execution_state",
        "//xla/ffi:ffi_api",
        "//xla/ffi/api:c_api",
        "//xla/hlo/ir:hlo",
        "//xla/runtime:object_pool",
        "//xla/service:buffer_assignment",
        "//xla/service:custom_call_status",
        "//xla/service:custom_call_status_internal",
        "//xla/service/gpu:buffer_allocations",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:device_memory_allocator",
        "//xla/stream_executor:stream",
        "//xla/tsl/platform:errors",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/memory",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/types:span",
    ],
)

xla_test(
    name = "custom_call_thunk_test",
    srcs = ["custom_call_thunk_test.cc"],
    backends = ["gpu"],
    deps = [
        ":custom_call_thunk",
        ":thunk",
        "//xla:executable_run_options",
        "//xla/service:custom_call_status_public_headers",
        "//xla/service:executable",
        "//xla/service:platform_util",
        "//xla/service/gpu:buffer_allocations",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:platform_manager",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_memory_allocator",
        "@com_google_absl//absl/status:statusor",
        "@com_google_googletest//:gtest_main",
        "@tsl//tsl/platform:status_matchers",
        "@tsl//tsl/platform:statusor",
        "@tsl//tsl/platform:test",
    ],
)

cc_library(
    name = "fft_thunk",
    srcs = ["fft_thunk.cc"],
    hdrs = ["fft_thunk.h"],
    deps = [
        ":thunk",
        "//xla:shape_util",
        "//xla:status_macros",
        "//xla:types",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/service:buffer_assignment",
        "//xla/stream_executor:blas",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:device_memory_allocator",
        "//xla/stream_executor:fft",
        "//xla/stream_executor:scratch_allocator",
        "//xla/stream_executor:stream_executor_h",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
        "@tsl//tsl/platform:logging",
        "@tsl//tsl/platform:statusor",
    ],
)

cc_library(
    name = "gemm_thunk",
    srcs = ["gemm_thunk.cc"],
    hdrs = ["gemm_thunk.h"],
    deps = [
        ":thunk",
        ":thunk_proto_cc",
        "//xla/service:buffer_assignment",
        "//xla/service/gpu:buffer_allocations",
        "//xla/service/gpu:matmul_utils",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:stream",
        "//xla/stream_executor/gpu:gpu_blas_lt",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/types:span",
        "@tsl//tsl/platform:logging",
        "@tsl//tsl/platform:statusor",
        "@tsl//tsl/profiler/lib:nvtx_utils",
    ],
)

xla_cc_test(
    name = "gemm_thunk_test",
    srcs = ["gemm_thunk_test.cc"],
    deps = [
        ":gemm_thunk",
        ":thunk",
        ":thunk_proto_cc",
        "//xla/service:buffer_assignment",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/platform:test",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_googletest//:gtest_main",
        "@tsl//tsl/platform:protobuf",
    ],
)

cc_library(
    name = "gpublas_lt_matmul_thunk",
    srcs = ["gpublas_lt_matmul_thunk.cc"],
    hdrs = ["gpublas_lt_matmul_thunk.h"],
    deps = [
        ":thunk",
        "//xla/hlo/ir:hlo",
        "//xla/service:buffer_assignment",
        "//xla/service/gpu:buffer_allocations",
        "//xla/service/gpu:matmul_utils",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:stream",
        "//xla/stream_executor/gpu:gpu_blas_lt",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
    ],
)

xla_test(
    name = "gpublas_lt_matmul_thunk_test",
    srcs = ["gpublas_lt_matmul_thunk_test.cc"],
    backends = ["gpu"],
    deps = [
        ":gpublas_lt_matmul_thunk",
        ":thunk",
        "//xla:error_spec",
        "//xla:executable_run_options",
        "//xla:shape_util",
        "//xla:xla_proto_cc",
        "//xla/hlo/ir:hlo",
        "//xla/service:buffer_assignment",
        "//xla/service:executable",
        "//xla/service/gpu:backend_configs_cc",
        "//xla/service/gpu:buffer_allocations",
        "//xla/service/gpu:cublas_cudnn",
        "//xla/service/gpu:matmul_utils",
        "//xla/service/gpu/transforms:gemm_rewriter",
        "//xla/stream_executor:device_description",
        "//xla/stream_executor:device_memory_allocator",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:semantic_version",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_memory_allocator",
        "//xla/stream_executor/gpu:gpu_blas_lt",
        "//xla/tests:hlo_test_base",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:env",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/time",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "infeed_thunk",
    srcs = ["infeed_thunk.cc"],
    hdrs = ["infeed_thunk.h"],
    deps = [
        ":thunk",
        "//xla:shape_tree",
        "//xla:shape_util",
        "//xla:status_macros",
        "//xla:util",
        "//xla/service/gpu:buffer_allocations",
        "//xla/service/gpu:gpu_transfer_manager",
        "//xla/service/gpu:io_feed_manager",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:device_memory_handle",
        "//xla/stream_executor:stream_executor_h",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@tsl//tsl/platform:errors",
    ],
)

cc_library(
    name = "kernel_thunk",
    srcs = ["kernel_thunk.cc"],
    hdrs = ["kernel_thunk.h"],
    deps = [
        ":thunk",
        ":thunk_proto_cc",
        "//xla:shape_util",
        "//xla:types",
        "//xla/codegen/emitters:kernel_arguments",
        "//xla/hlo/ir:hlo",
        "//xla/service:buffer_assignment",
        "//xla/service/gpu:launch_dimensions",
        "//xla/service/gpu:stream_executor_util",
        "//xla/service/gpu/kernels:custom_kernel",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:kernel",
        "//xla/stream_executor:launch_dim",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/stream_executor/gpu:tma_metadata",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
        "@llvm-project//llvm:Support",
    ],
)

xla_test(
    name = "kernel_thunk_test",
    srcs = ["kernel_thunk_test.cc"],
    backends = ["gpu"],
    deps = [
        ":command_buffer_cmd",
        ":command_buffer_cmd_emitter",
        ":command_buffer_thunk",
        ":kernel_thunk",
        ":sequential_thunk",
        ":thunk",
        ":thunk_proto_cc",
        "//xla:shape_util",
        "//xla:xla_data_proto_cc",
        "//xla/codegen/emitters:kernel_arguments",
        "//xla/service:buffer_assignment",
        "//xla/service:executable",
        "//xla/service/gpu:buffer_allocations",
        "//xla/service/gpu:launch_dimensions",
        "//xla/stream_executor:launch_dim",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:platform_manager",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/stream_executor/gpu:gpu_test_kernels",
        "//xla/stream_executor/gpu:tma_metadata",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/platform:test",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_googletest//:gtest_main",
        "@tsl//tsl/platform:protobuf",
    ],
)

cuda_library(
    name = "select_k_exec_raft",
    srcs = ["select_k_exec_raft.cc"],
    hdrs = ["select_k_exec.h"],
    copts = [
        "-fexceptions",
        "-DLIBCUDACXX_ENABLE_EXPERIMENTAL_MEMORY_RESOURCE",
    ],
    tags = ["cuda-only"],
    deps = [
        "//xla:status_macros",
        "//xla:types",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:device_memory_allocator",
        "//xla/stream_executor:scratch_allocator",
        "//xla/stream_executor:stream",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/strings",
        "@local_config_cuda//cuda:cuda_headers",
        "@raft//:raft_matrix",
    ],
)

cuda_library(
    name = "select_k_exec_stub",
    srcs = ["select_k_exec_stub.cc"],
    hdrs = ["select_k_exec.h"],
    deps = [
        "//xla:types",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:device_memory_allocator",
        "//xla/stream_executor:stream",
        "@com_google_absl//absl/status",
    ],
)

xla_test(
    name = "select_k_exec_raft_test",
    srcs = ["select_k_exec_raft_test.cc"],
    backends = [
        "gpu",
    ],
    tags = [
        "cuda-only",
    ],
    deps = [
        ":select_k_exec_raft",
        "//xla:types",
        "//xla:xla_data_proto_cc",
        "//xla:xla_proto_cc",
        "//xla/service:platform_util",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:platform_manager",
        "//xla/stream_executor:stream_executor_h",
        "//xla/stream_executor:stream_executor_memory_allocator",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/random",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/types:span",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "select_k_thunk",
    srcs = ["select_k_thunk.cc"],
    hdrs = ["select_k_thunk.h"],
    tags = ["gpu"],
    deps = [
        ":thunk",
        ":thunk_proto_cc",
        "//xla:shape_util",
        "//xla:types",
        "//xla/codegen/emitters:kernel_arguments",
        "//xla/hlo/ir:hlo",
        "//xla/service:buffer_assignment",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:device_memory_allocator",
        "//xla/stream_executor:stream",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/strings",
    ] + if_cuda_is_configured(
        [":select_k_exec_raft"],
        no_cuda = [":select_k_exec_stub"],
    ),
)

cc_library(
    name = "memset_thunk",
    srcs = ["memset_thunk.cc"],
    hdrs = ["memset_thunk.h"],
    deps = [
        ":thunk",
        "//xla/service:buffer_assignment",
        "//xla/stream_executor:device_memory",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/types:span",
    ],
)

xla_cc_test(
    name = "memset_thunk_test",
    srcs = ["memset_thunk_test.cc"],
    deps = [
        ":memset_thunk",
        ":thunk",
        ":thunk_proto_cc",
        "//xla/service:buffer_assignment",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_absl//absl/log:check",
        "@com_google_googletest//:gtest_main",
        "@tsl//tsl/platform:protobuf",
    ],
)

cc_library(
    name = "all_gather_thunk",
    srcs = ["all_gather_thunk.cc"],
    hdrs = ["all_gather_thunk.h"],
    deps = [
        ":collective_thunk",
        ":thunk",
        "//xla:shape_util",
        "//xla/backends/gpu/collectives:gpu_clique_key",
        "//xla/backends/gpu/collectives:gpu_collectives",
        "//xla/backends/gpu/collectives:gpu_communicator",
        "//xla/core/collectives:communicator",
        "//xla/hlo/ir:hlo",
        "//xla/service:collective_ops_utils",
        "//xla/service/gpu:backend_configs_cc",
        "//xla/service/gpu/transforms/collectives:collective_ops_utils",
        "//xla/stream_executor:stream",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:logging",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/types:span",
        "@tsl//tsl/platform:casts",
    ],
)

cc_library(
    name = "collective_kernel_thunk",
    srcs = ["collective_kernel_thunk.cc"],
    hdrs = ["collective_kernel_thunk.h"],
    deps = [
        ":all_reduce",
        ":collective_thunk",
        ":thunk",
        "//xla:shape_util",
        "//xla:status_macros",
        "//xla:xla_data_proto_cc",
        "//xla/backends/gpu/collectives:gpu_clique_key",
        "//xla/core/collectives:rank_id",
        "//xla/service:collective_ops_utils",
        "//xla/service:rendezvous",
        "//xla/service/gpu:launch_dimensions",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:device_memory_handle",
        "//xla/stream_executor:stream",
        "//xla/stream_executor/gpu:all_reduce_kernel",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
    ],
)

cc_library(
    name = "all_reduce_thunk",
    srcs = ["all_reduce_thunk.cc"],
    hdrs = ["all_reduce_thunk.h"],
    deps = [
        ":collective_kernel_thunk",
        ":collective_thunk",
        ":thunk",
        "//xla:status_macros",
        "//xla:xla_data_proto_cc",
        "//xla/backends/gpu/collectives:gpu_clique_key",
        "//xla/backends/gpu/collectives:gpu_collectives",
        "//xla/backends/gpu/collectives:gpu_communicator",
        "//xla/core/collectives:communicator",
        "//xla/hlo/ir:hlo",
        "//xla/service:collective_ops_utils",
        "//xla/service/gpu:backend_configs_cc",
        "//xla/service/gpu/transforms/collectives:collective_ops_utils",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:logging",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/types:span",
        "@tsl//tsl/platform:casts",
    ],
)

cc_library(
    name = "all_to_all_thunk",
    srcs = ["all_to_all_thunk.cc"],
    hdrs = ["all_to_all_thunk.h"],
    deps = [
        ":collective_thunk",
        ":thunk",
        "//xla:shape_util",
        "//xla:status_macros",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/backends/gpu/collectives:gpu_clique_key",
        "//xla/backends/gpu/collectives:gpu_collectives",
        "//xla/backends/gpu/collectives:gpu_communicator",
        "//xla/core/collectives:communicator",
        "//xla/core/collectives:rank_id",
        "//xla/hlo/ir:hlo",
        "//xla/service:collective_ops_utils",
        "//xla/service:rendezvous",
        "//xla/service/gpu/transforms/collectives:collective_ops_utils",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:memory_allocation",
        "//xla/stream_executor:stream",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:logging",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
    ],
)

cc_library(
    name = "ragged_all_to_all_thunk",
    srcs = ["ragged_all_to_all_thunk.cc"],
    hdrs = ["ragged_all_to_all_thunk.h"],
    tags = ["gpu"],
    deps = [
        ":collective_thunk",
        ":ragged_all_to_all",
        ":thunk",
        "//xla:shape_util",
        "//xla:xla_data_proto_cc",
        "//xla/backends/gpu/collectives:gpu_clique_key",
        "//xla/backends/gpu/collectives:gpu_collectives",
        "//xla/backends/gpu/collectives:gpu_communicator",
        "//xla/core/collectives:communicator",
        "//xla/core/collectives:rank_id",
        "//xla/hlo/ir:collective_op_group_mode",
        "//xla/hlo/ir:hlo",
        "//xla/service:rendezvous",
        "//xla/service/gpu/transforms/collectives:collective_ops_utils",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:device_memory_handle",
        "//xla/stream_executor:event",
        "//xla/stream_executor:memory_allocation",
        "//xla/stream_executor:stream",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/container:node_hash_map",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
        "@tsl//tsl/platform:casts",
    ],
)

cc_library(
    name = "collective_broadcast_thunk",
    srcs = ["collective_broadcast_thunk.cc"],
    hdrs = ["collective_broadcast_thunk.h"],
    deps = [
        ":collective_thunk",
        ":thunk",
        "//xla:xla_data_proto_cc",
        "//xla/backends/gpu/collectives:gpu_clique_key",
        "//xla/backends/gpu/collectives:gpu_collectives",
        "//xla/backends/gpu/collectives:gpu_communicator",
        "//xla/core/collectives:communicator",
        "//xla/core/collectives:rank_id",
        "//xla/hlo/ir:hlo",
        "//xla/service:collective_ops_utils",
        "//xla/service/gpu/transforms/collectives:collective_ops_utils",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:stream",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/types:span",
        "@tsl//tsl/platform:casts",
    ],
)

cc_library(
    name = "collective_permute_thunk",
    srcs = ["collective_permute_thunk.cc"],
    hdrs = ["collective_permute_thunk.h"],
    deps = [
        ":collective_thunk",
        ":p2p_thunk_common",
        ":thunk",
        "//xla:executable_run_options",
        "//xla:xla_data_proto_cc",
        "//xla/backends/gpu/collectives:gpu_clique_key",
        "//xla/backends/gpu/collectives:gpu_collectives",
        "//xla/backends/gpu/collectives:gpu_communicator",
        "//xla/core/collectives:communicator",
        "//xla/core/collectives:rank_id",
        "//xla/hlo/ir:hlo",
        "//xla/service:collective_ops_utils",
        "//xla/service:computation_placer",
        "//xla/service:global_device_id",
        "//xla/service:rendezvous",
        "//xla/service/gpu:backend_configs_cc",
        "//xla/service/gpu/transforms/collectives:collective_ops_utils",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:event",
        "//xla/stream_executor:stream",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:node_hash_map",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/time",
        "@tsl//tsl/platform:casts",
    ],
)

cc_library(
    name = "collective_thunk",
    srcs = ["collective_thunk.cc"],
    hdrs = ["collective_thunk.h"],
    deps = [
        ":thunk",
        "//xla:debug_options_flags",
        "//xla:shape_util",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla:xla_proto_cc",
        "//xla/backends/gpu/collectives:gpu_clique_key",
        "//xla/backends/gpu/collectives:gpu_collectives",
        "//xla/core/collectives:communicator",
        "//xla/core/collectives:rank_id",
        "//xla/hlo/ir:hlo",
        "//xla/hlo/translate/mhlo_to_hlo:attribute_exporter",
        "//xla/service:buffer_assignment",
        "//xla/service:collective_ops_utils",
        "//xla/service:computation_placer",
        "//xla/service:global_device_id",
        "//xla/service:rendezvous",
        "//xla/service/gpu:buffer_allocations",
        "//xla/service/llvm_ir:llvm_util",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:event",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/base",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:flat_hash_set",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/time",
        "@com_google_absl//absl/types:span",
        "@llvm-project//mlir:IR",
    ],
)

cc_library(
    name = "p2p_thunk_common",
    srcs = ["p2p_thunk_common.cc"],
    hdrs = ["p2p_thunk_common.h"],
    deps = [
        ":collective_thunk",
        "//xla:executable_run_options",
        "//xla:shape_util",
        "//xla:status_macros",
        "//xla:xla_data_proto_cc",
        "//xla/backends/gpu/collectives:gpu_clique_key",
        "//xla/hlo/ir:hlo",
        "//xla/hlo/parser:hlo_parser",
        "//xla/service:collective_ops_utils",
        "//xla/stream_executor:stream_executor_h",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/synchronization",
        "@llvm-project//mlir:IR",
        "@tsl//tsl/platform:statusor",
    ],
)

cc_library(
    name = "recv_thunk",
    srcs = ["recv_thunk.cc"],
    hdrs = ["recv_thunk.h"],
    deps = [
        ":collective_thunk",
        ":p2p_thunk_common",
        ":thunk",
        "//xla:status_macros",
        "//xla/backends/gpu/collectives:gpu_clique_key",
        "//xla/backends/gpu/collectives:gpu_collectives",
        "//xla/core/collectives:communicator",
        "//xla/core/collectives:rank_id",
        "//xla/hlo/ir:hlo",
        "//xla/service:collective_ops_utils",
        "//xla/service:computation_placer",
        "//xla/service:global_device_id",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:stream",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/strings:string_view",
    ],
)

cc_library(
    name = "send_thunk",
    srcs = ["send_thunk.cc"],
    hdrs = ["send_thunk.h"],
    deps = [
        ":collective_thunk",
        ":p2p_thunk_common",
        ":thunk",
        "//xla:status_macros",
        "//xla/backends/gpu/collectives:gpu_clique_key",
        "//xla/backends/gpu/collectives:gpu_collectives",
        "//xla/core/collectives:communicator",
        "//xla/core/collectives:rank_id",
        "//xla/hlo/ir:hlo",
        "//xla/service:collective_ops_utils",
        "//xla/service:computation_placer",
        "//xla/service:global_device_id",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:stream",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/strings:string_view",
    ],
)

cc_library(
    name = "collective_group_thunk",
    srcs = ["collective_group_thunk.cc"],
    hdrs = ["collective_group_thunk.h"],
    deps = [
        ":collective_thunk",
        ":thunk",
        "//xla:util",
        "//xla/backends/gpu/collectives:gpu_clique_key",
        "//xla/backends/gpu/collectives:gpu_communicator",
        "//xla/core/collectives:communicator",
        "//xla/hlo/ir:hlo",
        "//xla/stream_executor:event",
        "//xla/stream_executor:stream",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/container:flat_hash_set",
        "@com_google_absl//absl/functional:function_ref",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@tsl//tsl/platform:casts",
    ],
)

cc_library(
    name = "norm_thunk",
    srcs = ["norm_thunk.cc"],
    hdrs = ["norm_thunk.h"],
    deps = [
        ":thunk",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/service:buffer_assignment",
        "//xla/service/gpu:gpu_norm_runner",
        "//xla/stream_executor:stream_executor_h",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/synchronization",
    ],
)

cc_library(
    name = "outfeed_thunk",
    srcs = ["outfeed_thunk.cc"],
    hdrs = ["outfeed_thunk.h"],
    deps = [
        ":thunk",
        "//xla:shape_tree",
        "//xla:shape_util",
        "//xla:status_macros",
        "//xla:util",
        "//xla/service:buffer_assignment",
        "//xla/service/gpu:buffer_allocations",
        "//xla/service/gpu:gpu_transfer_manager",
        "//xla/service/gpu:io_feed_manager",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:stream_executor_h",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@tsl//tsl/platform:errors",
    ],
)

cc_library(
    name = "replica_id_thunk",
    srcs = ["replica_id_thunk.cc"],
    hdrs = ["replica_id_thunk.h"],
    deps = [
        ":thunk",
        ":thunk_proto_cc",
        "//xla/service:buffer_assignment",
        "//xla/service:computation_placer_hdr",
        "//xla/service:global_device_id",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/types:span",
    ],
)

cc_library(
    name = "sequential_thunk",
    srcs = ["sequential_thunk.cc"],
    hdrs = ["sequential_thunk.h"],
    deps = [
        ":annotation",
        ":thunk",
        ":thunk_proto_cc",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/functional:function_ref",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@tsl//tsl/profiler/lib:scoped_annotation",
        "@tsl//tsl/profiler/lib:traceme",
    ],
)

xla_cc_test(
    name = "sequential_thunk_test",
    srcs = ["sequential_thunk_test.cc"],
    deps = [
        ":sequential_thunk",
        ":thunk",
        ":thunk_proto_cc",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "host_send_recv_thunk",
    srcs = ["host_send_recv_thunk.cc"],
    hdrs = ["host_send_recv_thunk.h"],
    deps = [
        ":thunk",
        "//xla:shape_util",
        "//xla:xla_data_proto_cc",
        "//xla/service:buffer_assignment",
        "//xla/service:global_device_id",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:event",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/synchronization",
        "@tsl//tsl/profiler/lib:traceme",
    ],
)

cc_library(
    name = "thunk",
    srcs = ["thunk.cc"],
    hdrs = ["thunk.h"],
    deps = [
        ":thunk_proto_cc",
        "//xla:executable_run_options",
        "//xla:status_macros",
        "//xla:util",
        "//xla/backends/gpu/collectives:gpu_clique_key",
        "//xla/backends/gpu/collectives:gpu_cliques",
        "//xla/backends/gpu/collectives:gpu_collectives",
        "//xla/core/collectives:communicator",
        "//xla/core/collectives:rank_id",
        "//xla/ffi:execution_context",
        "//xla/hlo/ir:hlo",
        "//xla/service:buffer_assignment",
        "//xla/service:executable",
        "//xla/service:global_device_id",
        "//xla/service/gpu:backend_configs_cc",
        "//xla/service/gpu:buffer_allocations",
        "//xla/service/gpu:gpu_executable_run_options",
        "//xla/service/gpu:ir_emission_utils",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/lib/gtl:int_type",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base:nullability",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/functional:any_invocable",
        "@com_google_absl//absl/functional:function_ref",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/types:span",
    ],
)

xla_cc_test(
    name = "for_all_thunks_test",
    srcs = ["for_all_thunks_test.cc"],
    tags = ["gpu"],
    deps = [
        ":command_buffer_cmd",
        ":command_buffer_thunk",
        ":conditional_thunk",
        ":dynamic_slice_thunk",
        ":sequential_thunk",
        ":thunk",
        ":thunk_proto_cc",
        ":while_thunk",
        "//xla/service:buffer_assignment",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "triangular_solve_thunk",
    srcs = ["triangular_solve_thunk.cc"],
    hdrs = ["triangular_solve_thunk.h"],
    deps = [
        ":make_batch_pointers",
        ":thunk",
        ":thunk_proto_cc",
        "//xla:status_macros",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/service:buffer_assignment",
        "//xla/stream_executor:blas",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:stream",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/types:span",
    ],
)

xla_cc_test(
    name = "triangular_solve_thunk_test",
    srcs = ["triangular_solve_thunk_test.cc"],
    deps = [
        ":thunk",
        ":thunk_proto_cc",
        ":triangular_solve_thunk",
        "//xla/service:buffer_assignment",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_absl//absl/log:check",
        "@com_google_googletest//:gtest_main",
        "@tsl//tsl/platform:protobuf",
    ],
)

xla_cc_test(
    name = "replica_id_thunk_test",
    srcs = ["replica_id_thunk_test.cc"],
    deps = [
        ":replica_id_thunk",
        ":thunk",
        ":thunk_proto_cc",
        "//xla/service:buffer_assignment",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_absl//absl/log:check",
        "@com_google_googletest//:gtest_main",
        "@tsl//tsl/platform:protobuf",
    ],
)

cc_library(
    name = "while_thunk",
    srcs = ["while_thunk.cc"],
    hdrs = ["while_thunk.h"],
    deps = [
        ":host_memory_pool",
        ":sequential_thunk",
        ":thunk",
        ":thunk_proto_cc",
        "//xla:xla_data_proto_cc",
        "//xla/hlo/ir:hlo",
        "//xla/service:buffer_assignment",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/cleanup",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/functional:function_ref",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
        "@tsl//tsl/profiler/lib:traceme",
    ],
)

xla_test(
    name = "while_thunk_test",
    srcs = ["while_thunk_test.cc"],
    backends = ["gpu"],
    tags = ["test_migrated_to_hlo_runner_pjrt"],
    deps = [
        ":sequential_thunk",
        ":thunk",
        ":thunk_proto_cc",
        ":while_thunk",
        "//xla:executable_run_options",
        "//xla/hlo/ir:hlo",
        "//xla/hlo/testlib:verified_hlo_module",
        "//xla/service:buffer_assignment",
        "//xla/service:executable",
        "//xla/service:platform_util",
        "//xla/service/gpu:buffer_allocations",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:platform_manager",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_memory_allocator",
        "//xla/tests:hlo_pjrt_test_base",
        "//xla/tsl/platform:status_matchers",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/platform:test",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_googletest//:gtest_main",
        "@tsl//tsl/platform:protobuf",
    ],
)

cc_library(
    name = "wait_for_streams_thunk",
    srcs = ["wait_for_streams_thunk.cc"],
    hdrs = ["wait_for_streams_thunk.h"],
    deps = [
        ":thunk",
        ":thunk_proto_cc",
        "//xla/stream_executor:stream",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@tsl//tsl/platform:statusor",
    ],
)

xla_cc_test(
    name = "wait_for_streams_thunk_test",
    srcs = ["wait_for_streams_thunk_test.cc"],
    deps = [
        ":thunk",
        ":thunk_proto_cc",
        ":wait_for_streams_thunk",
        "//xla/tsl/platform:statusor",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "cudnn_thunk",
    srcs = ["cudnn_thunk.cc"],
    hdrs = ["cudnn_thunk.h"],
    deps = [
        ":thunk",
        ":thunk_proto_cc",
        "//xla/service:buffer_assignment",
        "//xla/service:buffer_assignment_proto_cc",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:dnn",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/types:span",
        "@tsl//tsl/profiler/lib:nvtx_utils",
    ],
)

xla_cc_test(
    name = "cudnn_thunk_test",
    srcs = ["cudnn_thunk_test.cc"],
    deps = [
        ":cudnn_thunk",
        ":thunk",
        ":thunk_proto_cc",
        "//xla/service:buffer_assignment",
        "//xla/tsl/platform:status_matchers",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_googletest//:gtest_main",
        "@tsl//tsl/platform:protobuf",
    ],
)

cc_library(
    name = "custom_call_target",
    hdrs = ["custom_call_target.h"],
    deps = [
        "//xla/service:custom_call_status",
    ],
)

cc_library(
    name = "host_memory_pool",
    srcs = ["host_memory_pool.cc"],
    hdrs = ["host_memory_pool.h"],
    deps = [
        "//xla:shape_util",
        "//xla:xla_data_proto_cc",
        "//xla/stream_executor:memory_allocation",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/memory",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/synchronization",
    ],
)

cc_library(
    name = "buffer_comparator",
    srcs = ["buffer_comparator.cc"],
    hdrs = ["buffer_comparator.h"],
    deps = [
        "//xla:shape_util",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/service/gpu:launch_dimensions",
        "//xla/stream_executor:device_description",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:device_memory_handle",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/stream_executor/gpu:buffer_comparator_kernel",
        "//xla/stream_executor/gpu:gpu_kernel_registry",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@eigen_archive//:eigen3",
    ],
)

xla_test(
    name = "buffer_comparator_test",
    srcs = ["buffer_comparator_test.cc"],
    backends = ["gpu"],
    disabled_backends = [],
    deps = [
        ":buffer_comparator",
        "//xla:shape_util",
        "//xla:types",
        "//xla:xla_data_proto_cc",
        "//xla/service:platform_util",
        "//xla/service/gpu:stream_executor_util",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:device_memory_handle",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:platform_manager",
        "//xla/stream_executor:stream",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/types:span",
        "@com_google_googletest//:gtest_main",
        "@tsl//tsl/platform:ml_dtypes",
        "@tsl//tsl/platform:test",
    ],
)

cc_library(
    name = "make_batch_pointers",
    srcs = ["make_batch_pointers.cc"],
    hdrs = ["make_batch_pointers.h"],
    deps = [
        "//xla:types",
        "//xla:util",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:launch_dim",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/stream_executor/gpu:gpu_kernel_registry",
        "//xla/stream_executor/gpu:make_batch_pointers_kernel",
        "//xla/stream_executor/rocm:rocm_platform_id",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/status",
    ],
)

xla_test(
    name = "make_batch_pointers_test",
    srcs = ["make_batch_pointers_test.cc"],
    backends = ["gpu"],
    deps = [
        ":make_batch_pointers",
        "//xla/service:platform_util",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/platform:status_matchers",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/status:statusor",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "all_reduce",
    srcs = ["all_reduce.cc"],
    hdrs = ["all_reduce.h"],
    deps = [
        "//xla:shape_util",
        "//xla:types",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/core/collectives:rank_id",
        "//xla/service:collective_ops_utils",
        "//xla/service/gpu:launch_dimensions",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:stream",
        "//xla/stream_executor/gpu:all_reduce_kernel",
        "//xla/stream_executor/gpu:gpu_kernel_registry",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/util:safe_reinterpret_cast",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/base",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/types:span",
    ],
)

xla_test(
    name = "all_reduce_test",
    srcs = ["all_reduce_test.cc"],
    backend_tags = {
        "gpu": [
            "multi_gpu_h100",
            "no_oss",
        ],
    },
    backends = ["gpu"],
    disabled_backends = [],
    deps = [
        ":all_reduce",
        "//xla:array",
        "//xla:literal",
        "//xla:literal_util",
        "//xla:shape_util",
        "//xla:status_macros",
        "//xla:types",
        "//xla:xla_data_proto_cc",
        "//xla/core/collectives:rank_id",
        "//xla/hlo/ir:hlo",
        "//xla/hlo/testlib:hlo_hardware_independent_test_base",
        "//xla/service:collective_ops_utils",
        "//xla/service:hlo_runner",
        "//xla/service:platform_util",
        "//xla/service/gpu:launch_dimensions",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:device_memory_handle",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:platform_manager",
        "//xla/stream_executor:stream",
        "//xla/stream_executor/gpu:all_reduce_kernel",
        "//xla/stream_executor/gpu:gpu_init",
        "//xla/stream_executor/host:host_platform",
        "//xla/tests:literal_test_util",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:status_matchers",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/platform:test",
        "//xla/tsl/platform:test_main",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/types:span",
        "@com_google_googletest//:gtest",
    ],
)

cc_library(
    name = "ragged_all_to_all",
    srcs = ["ragged_all_to_all.cc"],
    hdrs = ["ragged_all_to_all.h"],
    deps = [
        "//xla:shape_util",
        "//xla:types",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:launch_dim",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/stream_executor/gpu:gpu_kernel_registry",
        "//xla/stream_executor/gpu:ragged_all_to_all_kernel",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/types:span",
    ],
)

xla_test(
    name = "ragged_all_to_all_test",
    srcs = ["ragged_all_to_all_test.cc"],
    backends = ["gpu"],
    disabled_backends = [],
    deps = [
        ":ragged_all_to_all",
        "//xla:shape_util",
        "//xla:xla_data_proto_cc",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:device_memory_handle",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:platform_manager",
        "//xla/stream_executor:stream",
        "//xla/stream_executor/gpu:gpu_init",
        "//xla/stream_executor/host:host_platform",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:test",
        "//xla/tsl/platform:test_main",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/types:span",
        "@com_google_googletest//:gtest",
    ],
)

cc_library(
    name = "topk",
    srcs = ["topk.cc"],
    hdrs = ["topk.h"],
    compatible_with = [],
    tags = ["gpu"],
    deps = [
        "//xla:types",
        "//xla:xla_data_proto_cc",
        "//xla/service/gpu/kernels:custom_kernel",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:kernel",
        "//xla/stream_executor:kernel_spec",
        "//xla/stream_executor:launch_dim",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:platform_manager",
        "//xla/stream_executor/gpu:gpu_kernel_registry",
        "//xla/stream_executor/gpu:topk_kernel",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/numeric:bits",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
    ],
)

xla_test(
    name = "topk_test",
    srcs = ["topk_test.cc"],
    backends = ["gpu"],
    deps = [
        ":topk",
        "//xla:types",
        "//xla:xla_data_proto_cc",
        "//xla/service:platform_util",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:kernel",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:platform_manager",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/stream_executor/host:host_platform",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/platform:test",
        "//xla/tsl/platform:test_main",
        "@com_google_absl//absl/random",
        "@com_google_absl//absl/strings",
        "@com_google_googletest//:gtest",
    ],
)

tf_proto_library(
    name = "thunk_proto",
    srcs = [
        "thunk.proto",
    ],
    protodeps = [
        # keep sorted
        ":dynamic_slice_thunk_proto",
        "//xla:xla_data_proto",
        "//xla/service:buffer_assignment_proto",
        "//xla/service/gpu:launch_dimensions_proto",
        "//xla/stream_executor:launch_dim_proto",
        "//xla/stream_executor/gpu:gpu_blas_lt_proto",
        "//xla/stream_executor/gpu:tma_metadata_proto",
    ],
)

tf_proto_library(
    name = "dynamic_slice_thunk_proto",
    srcs = [
        "dynamic_slice_thunk.proto",
    ],
    protodeps = [
        # keep sorted
        "//xla:xla_data_proto",
        "//xla/service:buffer_assignment_proto",
        "//xla/service:hlo_proto",
    ],
)

cc_library(
    name = "thunk_proto_deserialization",
    srcs = ["thunk_proto_deserialization.cc"],
    hdrs = ["thunk_proto_deserialization.h"],
    deps = [
        ":conditional_thunk",
        ":copy_thunk",
        ":cudnn_thunk",
        ":gemm_thunk",
        ":kernel_thunk",
        ":memset_thunk",
        ":replica_id_thunk",
        ":sequential_thunk",
        ":thunk",
        ":thunk_proto_cc",
        ":triangular_solve_thunk",
        ":wait_for_streams_thunk",
        ":while_thunk",
        "//xla/service:buffer_assignment",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/types:span",
        "@tsl//tsl/platform:protobuf",
    ],
)

xla_cc_test(
    name = "thunk_proto_deserialization_test",
    srcs = ["thunk_proto_deserialization_test.cc"],
    deps = [
        ":conditional_thunk",
        ":copy_thunk",
        ":sequential_thunk",
        ":thunk",
        ":thunk_proto_cc",
        ":thunk_proto_deserialization",
        ":wait_for_streams_thunk",
        ":while_thunk",
        "//xla/service:buffer_assignment",
        "//xla/tsl/platform:status_matchers",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_googletest//:gtest_main",
        "@tsl//tsl/platform:protobuf",
    ],
)

cc_library(
    name = "nvshmem_collective_thunk",
    srcs = ["nvshmem_collective_thunk.cc"],
    hdrs = ["nvshmem_collective_thunk.h"],
    deps = [
        ":collective_thunk",
        ":thunk",
        "//xla:shape_util",
        "//xla:xla_data_proto_cc",
        "//xla/backends/gpu/collectives:gpu_clique_key",
        "//xla/backends/gpu/collectives:gpu_collectives",
        "//xla/core/collectives",
        "//xla/core/collectives:collectives_registry",
        "//xla/core/collectives:communicator",
        "//xla/service:computation_placer",
        "//xla/stream_executor:event",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "@com_google_absl//absl/base",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:flat_hash_set",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/synchronization",
        "@tsl//tsl/platform:errors",
        "@tsl//tsl/platform:logging",
        "@tsl//tsl/platform:statusor",
    ],
)

cc_library(
    name = "nvshmem_collective_permute_thunk",
    srcs = ["nvshmem_collective_permute_thunk.cc"],
    hdrs = ["nvshmem_collective_permute_thunk.h"],
    tags = ["gpu"],
    deps = [
        ":collective_permute_thunk",
        ":collective_thunk",
        ":nvshmem_collective_thunk",
        ":p2p_thunk_common",
        ":thunk",
        "//xla:xla_data_proto_cc",
        "//xla/backends/gpu/collectives:gpu_clique_key",
        "//xla/backends/gpu/collectives:gpu_collectives",
        "//xla/core/collectives:communicator",
        "//xla/core/collectives:rank_id",
        "//xla/hlo/ir:hlo",
        "//xla/service:collective_ops_utils",
        "//xla/service:computation_placer",
        "//xla/service:global_device_id",
        "//xla/service/gpu:backend_configs_cc",
        "//xla/service/gpu/transforms/collectives:collective_ops_utils",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:stream",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/strings:string_view",
    ],
)

cc_library(
    name = "nvshmem_send_thunk",
    srcs = ["nvshmem_send_thunk.cc"],
    hdrs = ["nvshmem_send_thunk.h"],
    tags = ["gpu"],
    deps = [
        ":collective_thunk",
        ":nvshmem_collective_thunk",
        ":p2p_thunk_common",
        ":thunk",
        "//xla:status_macros",
        "//xla/backends/gpu/collectives:gpu_collectives",
        "//xla/core/collectives:communicator",
        "//xla/core/collectives:rank_id",
        "//xla/hlo/ir:hlo",
        "//xla/runtime:device_id",
        "//xla/service:collective_ops_utils",
        "//xla/service:computation_placer",
        "//xla/service/gpu:backend_configs_cc",
        "//xla/service/gpu/transforms/collectives:collective_ops_utils",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:stream",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
    ],
)

cc_library(
    name = "nvshmem_recv_thunk",
    srcs = ["nvshmem_recv_thunk.cc"],
    hdrs = ["nvshmem_recv_thunk.h"],
    tags = ["gpu"],
    deps = [
        ":collective_thunk",
        ":nvshmem_collective_thunk",
        ":p2p_thunk_common",
        ":thunk",
        "//xla:status_macros",
        "//xla/backends/gpu/collectives:gpu_collectives",
        "//xla/core/collectives:communicator",
        "//xla/core/collectives:rank_id",
        "//xla/hlo/ir:hlo",
        "//xla/runtime:device_id",
        "//xla/service:collective_ops_utils",
        "//xla/service:computation_placer",
        "//xla/service/gpu:backend_configs_cc",
        "//xla/service/gpu/transforms/collectives:collective_ops_utils",
        "//xla/stream_executor:stream",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/strings:str_format",
    ],
)

cc_library(
    name = "nvshmem_all_reduce_thunk",
    srcs = ["nvshmem_all_reduce_thunk.cc"],
    hdrs = ["nvshmem_all_reduce_thunk.h"],
    deps = [
        ":all_reduce_thunk",
        ":collective_thunk",
        ":nvshmem_collective_thunk",
        ":thunk",
        "//xla:xla_data_proto_cc",
        "//xla/backends/gpu/collectives:gpu_collectives",
        "//xla/core/collectives:communicator",
        "//xla/hlo/ir:hlo",
        "//xla/service:collective_ops_utils",
        "//xla/service:computation_placer",
        "//xla/service/gpu:backend_configs_cc",
        "//xla/service/gpu/transforms/collectives:collective_ops_utils",
        "//xla/stream_executor:stream",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/types:span",
    ],
)

cc_library(
    name = "host_execute_thunk",
    srcs = ["host_execute_thunk.cc"],
    hdrs = ["host_execute_thunk.h"],
    deps = [
        ":thunk",
        ":thunk_proto_cc",
        "//xla:executable_run_options",
        "//xla:shape_tree",
        "//xla:shape_util",
        "//xla:types",
        "//xla:util",
        "//xla/backends/gpu/host_offloading:gpu_host_offloading_allocator",
        "//xla/core/host_offloading:host_offloading_allocator",
        "//xla/core/host_offloading:host_offloading_buffer",
        "//xla/core/host_offloading:host_offloading_executable",
        "//xla/core/host_offloading:host_offloading_executable_proto_cc",
        "//xla/core/host_offloading:host_offloading_nanort_executable",
        "//xla/hlo/ir:hlo",
        "//xla/service:buffer_assignment",
        "//xla/service/gpu:buffer_allocations",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/framework:allocator",
        "//xla/tsl/platform:env",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
        "@tsl//tsl/platform:platform_port",
        "@tsl//tsl/profiler/lib:traceme",
    ],
)

xla_test(
    name = "host_execute_thunk_test",
    srcs = ["host_execute_thunk_test.cc"],
    backends = ["gpu"],
    deps = [
        ":host_execute_thunk",
        ":thunk",
        "//xla:executable_run_options",
        "//xla:literal",
        "//xla:literal_util",
        "//xla:shape_util",
        "//xla:util",
        "//xla/backends/cpu:alignment",
        "//xla/core/host_offloading:host_offloading_executable",
        "//xla/hlo/parser:hlo_parser",
        "//xla/service:buffer_assignment",
        "//xla/service:executable",
        "//xla/service:platform_util",
        "//xla/service/gpu:buffer_allocations",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:platform_manager",
        "//xla/stream_executor:stream_executor_h",
        "//xla/stream_executor:stream_executor_memory_allocator",
        "//xla/tests:literal_test_util",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:status_matchers",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/strings",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "thunk_pass_pipeline",
    srcs = ["thunk_pass_pipeline.cc"],
    hdrs = ["thunk_pass_pipeline.h"],
    deps = [
        "//xla/backends/gpu/runtime:sequential_thunk",
        "//xla/service:hlo_module_config",
        "//xla/stream_executor:device_description",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
    ],
)

xla_cc_test(
    name = "thunk_pass_pipeline_test",
    srcs = ["thunk_pass_pipeline_test.cc"],
    deps = [
        ":thunk_pass_pipeline",
        "//xla/backends/gpu/runtime:sequential_thunk",
        "//xla/backends/gpu/runtime:thunk",
        "//xla/service:hlo_module_config",
        "//xla/stream_executor:device_description",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "command_buffer_conversion_pass",
    srcs = ["command_buffer_conversion_pass.cc"],
    hdrs = ["command_buffer_conversion_pass.h"],
    deps = [
        ":thunk_pass_pipeline",
        "//xla:util",
        "//xla/backends/gpu/runtime:command_buffer_cmd",
        "//xla/backends/gpu/runtime:command_buffer_cmd_emitter",
        "//xla/backends/gpu/runtime:command_buffer_thunk",
        "//xla/backends/gpu/runtime:conditional_thunk",
        "//xla/backends/gpu/runtime:copy_thunk",
        "//xla/backends/gpu/runtime:custom_call_thunk",
        "//xla/backends/gpu/runtime:sequential_thunk",
        "//xla/backends/gpu/runtime:thunk",
        "//xla/backends/gpu/runtime:while_thunk",
        "//xla/ffi:ffi_api",
        "//xla/stream_executor:device_description",
        "//xla/stream_executor:semantic_version",
        "//xla/stream_executor/cuda:cuda_compute_capability",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/container:flat_hash_set",
        "@com_google_absl//absl/functional:overload",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/types:span",
        "@tsl//tsl/platform",
        "@tsl//tsl/profiler/lib:traceme",
    ],
)

xla_test(
    name = "command_buffer_conversion_pass_test",
    srcs = ["command_buffer_conversion_pass_test.cc"],
    backends = ["gpu"],
    deps = [
        ":command_buffer_conversion_pass",
        "//xla:shape_util",
        "//xla/backends/gpu/collectives:gpu_clique_key",
        "//xla/backends/gpu/runtime:all_gather_thunk",
        "//xla/backends/gpu/runtime:collective_thunk",
        "//xla/backends/gpu/runtime:command_buffer_thunk",
        "//xla/backends/gpu/runtime:conditional_thunk",
        "//xla/backends/gpu/runtime:copy_thunk",
        "//xla/backends/gpu/runtime:cudnn_thunk",
        "//xla/backends/gpu/runtime:custom_call_thunk",
        "//xla/backends/gpu/runtime:gemm_thunk",
        "//xla/backends/gpu/runtime:replica_id_thunk",
        "//xla/backends/gpu/runtime:sequential_thunk",
        "//xla/backends/gpu/runtime:thunk",
        "//xla/backends/gpu/runtime:while_thunk",
        "//xla/hlo/ir:hlo",
        "//xla/service:buffer_assignment",
        "//xla/service:hlo_module_config",
        "//xla/service:platform_util",
        "//xla/service/gpu:gpu_device_info_for_tests",
        "//xla/service/gpu:matmul_utils",
        "//xla/stream_executor:blas",
        "//xla/stream_executor:device_description",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:platform_manager",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/platform:status",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/status:status_matchers",
        "@com_google_absl//absl/strings",
        "@com_google_googletest//:gtest_main",
    ],
)
